//===- FIRRTLTypes.cpp - Implement the FIRRTL dialect type system ---------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This file implement the FIRRTL dialect type system.
//
//===----------------------------------------------------------------------===//

#include "circt/Dialect/FIRRTL/FIRRTLTypes.h"
#include "circt/Dialect/FIRRTL/FIRRTLOps.h"
#include "mlir/IR/DialectImplementation.h"
#include "llvm/ADT/StringExtras.h"
#include "llvm/ADT/StringSwitch.h"
#include "llvm/ADT/TypeSwitch.h"

using namespace circt;
using namespace firrtl;

using mlir::OptionalParseResult;
using mlir::TypeStorageAllocator;

//===----------------------------------------------------------------------===//
// TableGen generated logic.
//===----------------------------------------------------------------------===//

// Provide the autogenerated implementation for types.
#define GET_TYPEDEF_CLASSES
#include "circt/Dialect/FIRRTL/FIRRTLTypes.cpp.inc"

//===----------------------------------------------------------------------===//
// Type Printing
//===----------------------------------------------------------------------===//

/// Print a type with a custom printer implementation.
///
/// This only prints a subset of all types in the dialect. Use `printNestedType`
/// instead, which will call this function in turn, as appropriate.
static LogicalResult customTypePrinter(Type type, AsmPrinter &os) {
  auto printWidthQualifier = [&](std::optional<int32_t> width) {
    if (width)
      os << '<' << *width << '>';
  };
  bool anyFailed = false;
  TypeSwitch<Type>(type)
      .Case<ClockType>([&](auto) { os << "clock"; })
      .Case<ResetType>([&](auto) { os << "reset"; })
      .Case<AsyncResetType>([&](auto) { os << "asyncreset"; })
      .Case<SIntType>([&](auto sIntType) {
        os << "sint";
        printWidthQualifier(sIntType.getWidth());
      })
      .Case<UIntType>([&](auto uIntType) {
        os << "uint";
        printWidthQualifier(uIntType.getWidth());
      })
      .Case<AnalogType>([&](auto analogType) {
        os << "analog";
        printWidthQualifier(analogType.getWidth());
      })
      .Case<BundleType>([&](auto bundleType) {
        os << "bundle<";
        llvm::interleaveComma(bundleType, os,
                              [&](BundleType::BundleElement element) {
                                os << element.name.getValue();
                                if (element.isFlip)
                                  os << " flip";
                                os << ": ";
                                printNestedType(element.type, os);
                              });
        os << '>';
      })
      .Case<FVectorType>([&](auto vectorType) {
        os << "vector<";
        printNestedType(vectorType.getElementType(), os);
        os << ", " << vectorType.getNumElements() << '>';
      })
      .Case<RefType>([&](auto refType) {
        os << "ref<";
        printNestedType(refType.getType(), os);
        os << '>';
      })
      .Default([&](auto) { anyFailed = true; });
  return failure(anyFailed);
}

/// Print a type defined by this dialect.
void circt::firrtl::printNestedType(Type type, AsmPrinter &os) {
  // Try the custom type printer.
  if (succeeded(customTypePrinter(type, os)))
    return;

  // None of the above recognized the type, so we bail.
  assert(false && "type to print unknown to FIRRTL dialect");
}

//===----------------------------------------------------------------------===//
// Type Parsing
//===----------------------------------------------------------------------===//

/// Parse a type with a custom parser implementation.
///
/// This only accepts a subset of all types in the dialect. Use `parseType`
/// instead, which will call this function in turn, as appropriate.
///
/// Returns `std::nullopt` if the type `name` is not covered by the custom
/// parsers. Otherwise returns success or failure as appropriate. On success,
/// `result` is set to the resulting type.
///
/// ```plain
/// firrtl-type
///   ::= clock
///   ::= reset
///   ::= asyncreset
///   ::= sint ('<' int '>')?
///   ::= uint ('<' int '>')?
///   ::= analog ('<' int '>')?
///   ::= bundle '<' (bundle-elt (',' bundle-elt)*)? '>'
///   ::= vector '<' type ',' int '>'
///
/// bundle-elt ::= identifier ':' type
/// ```
static OptionalParseResult customTypeParser(AsmParser &parser, StringRef name,
                                            Type &result) {
  auto *context = parser.getContext();
  if (name.equals("clock"))
    return result = ClockType::get(context), success();
  if (name.equals("reset"))
    return result = ResetType::get(context), success();
  if (name.equals("asyncreset"))
    return result = AsyncResetType::get(context), success();

  if (name.equals("sint") || name.equals("uint") || name.equals("analog")) {
    // Parse the width specifier if it exists.
    int32_t width = -1;
    if (!parser.parseOptionalLess()) {
      if (parser.parseInteger(width) || parser.parseGreater())
        return failure();

      if (width < 0)
        return parser.emitError(parser.getNameLoc(), "unknown width"),
               failure();
    }

    if (name.equals("sint"))
      result = SIntType::get(context, width);
    else if (name.equals("uint"))
      result = UIntType::get(context, width);
    else {
      assert(name.equals("analog"));
      result = AnalogType::get(context, width);
    }
    return success();
  }

  if (name.equals("bundle")) {
    SmallVector<BundleType::BundleElement, 4> elements;

    auto parseBundleElement = [&]() -> ParseResult {
      std::string nameStr;
      StringRef name;
      FIRRTLBaseType type;

      // The 'name' can be an identifier or an integer.
      uint32_t fieldIntName;
      auto intName = parser.parseOptionalInteger(fieldIntName);
      if (intName.has_value()) {
        if (failed(intName.value()))
          return failure();
        nameStr = llvm::utostr(fieldIntName);
        name = nameStr;
      } else {
        // Otherwise must be an identifier.
        if (parser.parseKeyword(&name))
          return failure();
      }

      bool isFlip = succeeded(parser.parseOptionalKeyword("flip"));
      if (parser.parseColon() || parseNestedBaseType(type, parser))
        return failure();

      elements.push_back({StringAttr::get(context, name), isFlip, type});
      return success();
    };

    if (parser.parseCommaSeparatedList(mlir::AsmParser::Delimiter::LessGreater,
                                       parseBundleElement))
      return failure();

    return result = BundleType::get(context, elements), success();
  }

  if (name.equals("vector")) {
    FIRRTLBaseType elementType;
    uint64_t width = 0;

    if (parser.parseLess() || parseNestedBaseType(elementType, parser) ||
        parser.parseComma() || parser.parseInteger(width) ||
        parser.parseGreater())
      return failure();

    return result = FVectorType::get(elementType, width), success();
  }

  if (name.equals("ref")) {
    FIRRTLBaseType type;
    if (parser.parseLess() || parseNestedBaseType(type, parser) ||
        parser.parseGreater())
      return failure();

    if (failed(RefType::verify(
            [&]() { return parser.emitError(parser.getNameLoc()); }, type)))
      return failure();

    return result = RefType::get(type), success();
  }

  return {};
}

/// Parse a type defined by this dialect.
///
/// This will first try the generated type parsers and then resort to the custom
/// parser implementation. Emits an error and returns failure if `name` does not
/// refer to a type defined in this dialect.
static ParseResult parseType(Type &result, StringRef name, AsmParser &parser) {
  // Try the custom type parser.
  OptionalParseResult parseResult = customTypeParser(parser, name, result);
  if (parseResult.has_value())
    return parseResult.value();

  // None of the above recognized the type, so we bail.
  parser.emitError(parser.getNameLoc(), "unknown FIRRTL dialect type: \"")
      << name << "\"";
  return failure();
}

/// Parse a `FIRRTLType` with a `name` that has already been parsed.
///
/// Note that only a subset of types defined in the FIRRTL dialect inherit from
/// `FIRRTLType`. Use `parseType` to parse *any* of the defined types.
static ParseResult parseFIRRTLType(FIRRTLType &result, StringRef name,
                                   AsmParser &parser) {
  Type type;
  if (failed(parseType(type, name, parser)))
    return failure();
  result = type.dyn_cast<FIRRTLType>();
  if (result)
    return success();
  parser.emitError(parser.getNameLoc(), "unknown FIRRTL type: \"")
      << name << "\"";
  return failure();
}

static ParseResult parseFIRRTLBaseType(FIRRTLBaseType &result, StringRef name,
                                       AsmParser &parser) {
  FIRRTLType type;
  if (failed(parseFIRRTLType(type, name, parser)))
    return failure();
  if (auto base = type.dyn_cast<FIRRTLBaseType>()) {
    result = base;
    return success();
  }
  parser.emitError(parser.getNameLoc(), "expected base type, found ") << type;
  return failure();
}

/// Parse a `FIRRTLType`.
///
/// Note that only a subset of types defined in the FIRRTL dialect inherit from
/// `FIRRTLType`. Use `parseType` to parse *any* of the defined types.
ParseResult circt::firrtl::parseNestedType(FIRRTLType &result,
                                           AsmParser &parser) {
  StringRef name;
  if (parser.parseKeyword(&name))
    return failure();
  return parseFIRRTLType(result, name, parser);
}

ParseResult circt::firrtl::parseNestedBaseType(FIRRTLBaseType &result,
                                               AsmParser &parser) {
  StringRef name;
  if (parser.parseKeyword(&name))
    return failure();
  return parseFIRRTLBaseType(result, name, parser);
}

//===---------------------------------------------------------------------===//
// Dialect Type Parsing and Printing
//===----------------------------------------------------------------------===//

/// Print a type registered to this dialect.
void FIRRTLDialect::printType(Type type, DialectAsmPrinter &os) const {
  printNestedType(type, os);
}

/// Parse a type registered to this dialect.
Type FIRRTLDialect::parseType(DialectAsmParser &parser) const {
  StringRef name;
  Type result;
  if (parser.parseKeyword(&name) || ::parseType(result, name, parser))
    return Type();
  return result;
}

//===----------------------------------------------------------------------===//
// Recursive Type Properties
//===----------------------------------------------------------------------===//

enum {
  /// Bit set if the type only contains passive elements.
  IsPassiveBitMask = 0x1,
  /// Bit set if the type contains an analog type.
  ContainsAnalogBitMask = 0x2,
  /// Bit set fi the type has any uninferred bit widths.
  HasUninferredWidthBitMask = 0x4,
};

/// Unpack `RecursiveTypeProperties` from a bunch of bits.
RecursiveTypeProperties RecursiveTypeProperties::fromFlags(unsigned flags) {
  return RecursiveTypeProperties{
      (flags & IsPassiveBitMask) != 0,
      (flags & ContainsAnalogBitMask) != 0,
      (flags & HasUninferredWidthBitMask) != 0,
  };
}

/// Pack `RecursiveTypeProperties` as a bunch of bits.
unsigned RecursiveTypeProperties::toFlags() const {
  unsigned flags = 0;
  if (isPassive)
    flags |= IsPassiveBitMask;
  if (containsAnalog)
    flags |= ContainsAnalogBitMask;
  if (hasUninferredWidth)
    flags |= HasUninferredWidthBitMask;
  return flags;
}

//===----------------------------------------------------------------------===//
// FIRRTLBaseType Implementation
//===----------------------------------------------------------------------===//

/// Return true if this is a 'ground' type, aka a non-aggregate type.
bool FIRRTLBaseType::isGround() {
  return TypeSwitch<FIRRTLBaseType, bool>(*this)
      .Case<ClockType, ResetType, AsyncResetType, SIntType, UIntType,
            AnalogType>([](Type) { return true; })
      .Case<BundleType, FVectorType>([](Type) { return false; })
      .Default([](Type) {
        llvm_unreachable("unknown FIRRTL type");
        return false;
      });
}

/// Return a pair with the 'isPassive' and 'containsAnalog' bits.
RecursiveTypeProperties FIRRTLBaseType::getRecursiveTypeProperties() {
  return TypeSwitch<FIRRTLBaseType, RecursiveTypeProperties>(*this)
      .Case<ClockType, ResetType, AsyncResetType>([](Type) {
        return RecursiveTypeProperties{true, false, false};
      })
      .Case<SIntType, UIntType>([](auto type) {
        return RecursiveTypeProperties{true, false, !type.hasWidth()};
      })
      .Case<AnalogType>([](auto type) {
        return RecursiveTypeProperties{true, true, !type.hasWidth()};
      })
      .Case<BundleType>([](BundleType bundleType) {
        return bundleType.getRecursiveTypeProperties();
      })
      .Case<FVectorType>([](FVectorType vectorType) {
        return vectorType.getRecursiveTypeProperties();
      })
      .Default([](Type) {
        llvm_unreachable("unknown FIRRTL type");
        return RecursiveTypeProperties{};
      });
}

/// Return this type with any flip types recursively removed from itself.
FIRRTLBaseType FIRRTLBaseType::getPassiveType() {
  return TypeSwitch<FIRRTLBaseType, FIRRTLBaseType>(*this)
      .Case<ClockType, ResetType, AsyncResetType, SIntType, UIntType,
            AnalogType>([&](Type) { return *this; })
      .Case<BundleType>(
          [](BundleType bundleType) { return bundleType.getPassiveType(); })
      .Case<FVectorType>(
          [](FVectorType vectorType) { return vectorType.getPassiveType(); })
      .Default([](Type) {
        llvm_unreachable("unknown FIRRTL type");
        return FIRRTLBaseType();
      });
}

/// Return this type with all ground types replaced with UInt<1>.  This is
/// used for `mem` operations.
FIRRTLBaseType FIRRTLBaseType::getMaskType() {
  return TypeSwitch<FIRRTLBaseType, FIRRTLBaseType>(*this)
      .Case<ClockType, ResetType, AsyncResetType, SIntType, UIntType,
            AnalogType>(
          [&](Type) { return UIntType::get(this->getContext(), 1); })
      .Case<BundleType>([&](BundleType bundleType) {
        SmallVector<BundleType::BundleElement, 4> newElements;
        newElements.reserve(bundleType.getElements().size());
        for (auto elt : bundleType)
          newElements.push_back(
              {elt.name, false /* FIXME */, elt.type.getMaskType()});
        return BundleType::get(this->getContext(), newElements);
      })
      .Case<FVectorType>([](FVectorType vectorType) {
        return FVectorType::get(vectorType.getElementType().getMaskType(),
                                vectorType.getNumElements());
      })
      .Default([](Type) {
        llvm_unreachable("unknown FIRRTL type");
        return FIRRTLBaseType();
      });
}

/// Remove the widths from this type. All widths are replaced with an
/// unknown width.
FIRRTLBaseType FIRRTLBaseType::getWidthlessType() {
  return TypeSwitch<FIRRTLBaseType, FIRRTLBaseType>(*this)
      .Case<ClockType, ResetType, AsyncResetType>([](auto a) { return a; })
      .Case<UIntType, SIntType, AnalogType>(
          [&](auto a) { return a.get(this->getContext(), -1); })
      .Case<BundleType>([&](auto a) {
        SmallVector<BundleType::BundleElement, 4> newElements;
        newElements.reserve(a.getElements().size());
        for (auto elt : a)
          newElements.push_back(
              {elt.name, elt.isFlip, elt.type.getWidthlessType()});
        return BundleType::get(this->getContext(), newElements);
      })
      .Case<FVectorType>([](auto a) {
        return FVectorType::get(a.getElementType().getWidthlessType(),
                                a.getNumElements());
      })
      .Default([](auto) {
        llvm_unreachable("unknown FIRRTL type");
        return FIRRTLBaseType();
      });
}

/// If this is an IntType, AnalogType, or sugar type for a single bit (Clock,
/// Reset, etc) then return the bitwidth.  Return -1 if the is one of these
/// types but without a specified bitwidth.  Return -2 if this isn't a simple
/// type.
int32_t FIRRTLBaseType::getBitWidthOrSentinel() {
  return TypeSwitch<FIRRTLBaseType, int32_t>(*this)
      .Case<ClockType, ResetType, AsyncResetType>([](Type) { return 1; })
      .Case<SIntType, UIntType>(
          [&](IntType intType) { return intType.getWidthOrSentinel(); })
      .Case<AnalogType>(
          [](AnalogType analogType) { return analogType.getWidthOrSentinel(); })
      .Case<BundleType, FVectorType>([](Type) { return -2; })
      .Default([](Type) {
        llvm_unreachable("unknown FIRRTL type");
        return -2;
      });
}

/// Return true if this is a type usable as a reset. This must be
/// either an abstract reset, a concrete 1-bit UInt, an
/// asynchronous reset, or an uninfered width UInt.
bool FIRRTLBaseType::isResetType() {
  return TypeSwitch<FIRRTLBaseType, bool>(*this)
      .Case<ResetType, AsyncResetType>([](Type) { return true; })
      .Case<UIntType>(
          [](UIntType a) { return !a.hasWidth() || a.getWidth() == 1; })
      .Default([](Type) { return false; });
}

uint64_t FIRRTLBaseType::getMaxFieldID() {
  return TypeSwitch<FIRRTLBaseType, uint64_t>(*this)
      .Case<AnalogType, ClockType, ResetType, AsyncResetType, SIntType,
            UIntType>([](Type) { return 0; })
      .Case<BundleType, FVectorType>(
          [](auto type) { return type.getMaxFieldID(); })
      .Default([](Type) {
        llvm_unreachable("unknown FIRRTL type");
        return -1;
      });
}

std::pair<FIRRTLBaseType, uint64_t>
FIRRTLBaseType::getSubTypeByFieldID(uint64_t fieldID) {
  return TypeSwitch<FIRRTLBaseType, std::pair<FIRRTLBaseType, uint64_t>>(*this)
      .Case<AnalogType, ClockType, ResetType, AsyncResetType, SIntType,
            UIntType>([&](FIRRTLBaseType t) {
        assert(!fieldID && "non-aggregate types must have a field id of 0");
        return std::pair(t, 0);
      })
      .Case<BundleType, FVectorType>(
          [&](auto type) { return type.getSubTypeByFieldID(fieldID); })
      .Default([](Type) {
        llvm_unreachable("unknown FIRRTL type");
        return std::pair(FIRRTLBaseType(), 0);
      });
}

FIRRTLBaseType FIRRTLBaseType::getFinalTypeByFieldID(uint64_t fieldID) {
  std::pair<FIRRTLBaseType, uint64_t> pair(*this, fieldID);
  while (pair.second)
    pair = pair.first.getSubTypeByFieldID(pair.second);
  return pair.first;
}

std::pair<uint64_t, bool> FIRRTLBaseType::rootChildFieldID(uint64_t fieldID,
                                                           uint64_t index) {
  return TypeSwitch<FIRRTLBaseType, std::pair<uint64_t, bool>>(*this)
      .Case<AnalogType, ClockType, ResetType, AsyncResetType, SIntType,
            UIntType>([&](Type) { return std::make_pair(0, fieldID == 0); })
      .Case<BundleType, FVectorType>(
          [&](auto type) { return type.rootChildFieldID(fieldID, index); })
      .Default([](Type) {
        llvm_unreachable("unknown FIRRTL type");
        return std::make_pair(0, false);
      });
}

uint64_t FIRRTLBaseType::getGroundFields() const {
  return TypeSwitch<FIRRTLBaseType, uint64_t>(*this)
      .Case<BundleType>([](auto type) {
        unsigned sum = 0;
        for (auto &field : type.getElements())
          sum += field.type.getGroundFields();
        return sum;
      })
      .Case<FVectorType>([](auto type) {
        return type.getNumElements() * type.getElementType().getGroundFields();
      })
      .Default([](Type) { return 1; });
}

/// Helper to implement the equivalence logic for a pair of bundle elements.
/// Note that the FIRRTL spec requires bundle elements to have the same
/// orientation, but this only compares their passive types. The FIRRTL dialect
/// differs from the spec in how it uses flip types for module output ports and
/// canonicalizes flips in bundles, so only passive types can be compared here.
static bool areBundleElementsEquivalent(BundleType::BundleElement destElement,
                                        BundleType::BundleElement srcElement) {
  if (destElement.name != srcElement.name)
    return false;
  if (destElement.isFlip != srcElement.isFlip)
    return false;

  return areTypesEquivalent(destElement.type, srcElement.type);
}

/// Returns whether the two types are equivalent.  This implements the exact
/// definition of type equivalence in the FIRRTL spec.  If the types being
/// compared have any outer flips that encode FIRRTL module directions (input or
/// output), these should be stripped before using this method.
bool firrtl::areTypesEquivalent(FIRRTLType destFType, FIRRTLType srcFType) {
  auto destType = destFType.dyn_cast<FIRRTLBaseType>();
  auto srcType = srcFType.dyn_cast<FIRRTLBaseType>();

  // For non-base types, only equivalent if identical.
  if (!destType || !srcType)
    return destFType == srcFType;

  // Reset types can be driven by UInt<1>, AsyncReset, or Reset types.
  if (destType.isa<ResetType>())
    return srcType.isResetType();

  // Reset types can drive UInt<1>, AsyncReset, or Reset types.
  if (srcType.isa<ResetType>())
    return destType.isResetType();

  // Vector types can be connected if they have the same size and element type.
  auto destVectorType = destType.dyn_cast<FVectorType>();
  auto srcVectorType = srcType.dyn_cast<FVectorType>();
  if (destVectorType && srcVectorType)
    return destVectorType.getNumElements() == srcVectorType.getNumElements() &&
           areTypesEquivalent(destVectorType.getElementType(),
                              srcVectorType.getElementType());

  // Bundle types can be connected if they have the same size, element names,
  // and element types.
  auto destBundleType = destType.dyn_cast<BundleType>();
  auto srcBundleType = srcType.dyn_cast<BundleType>();
  if (destBundleType && srcBundleType) {
    auto destElements = destBundleType.getElements();
    auto srcElements = srcBundleType.getElements();
    size_t numDestElements = destElements.size();
    if (numDestElements != srcElements.size())
      return false;

    for (size_t i = 0; i < numDestElements; ++i) {
      auto destElement = destElements[i];
      auto srcElement = srcElements[i];
      if (!areBundleElementsEquivalent(destElement, srcElement))
        return false;
    }
    return true;
  }

  // Ground types can be connected if their passive, widthless versions
  // are equal.
  return destType.getWidthlessType() == srcType.getWidthlessType();
}

/// Returns whether the two types are weakly equivalent.
bool firrtl::areTypesWeaklyEquivalent(FIRRTLType destFType, FIRRTLType srcFType,
                                      bool destFlip, bool srcFlip) {
  auto destType = destFType.dyn_cast<FIRRTLBaseType>();
  auto srcType = srcFType.dyn_cast<FIRRTLBaseType>();

  // For non-base types, only equivalent if identical.
  if (!destType || !srcType)
    return destFType == srcFType;

  // Reset types can be driven by UInt<1>, AsyncReset, or Reset types.
  if (destType.isa<ResetType>())
    return srcType.isResetType();

  // Reset types can drive UInt<1>, AsyncReset, or Reset types.
  if (srcType.isa<ResetType>())
    return destType.isResetType();

  // Vector types can be connected if their element types are weakly equivalent.
  // Size doesn't matter.
  auto destVectorType = destType.dyn_cast<FVectorType>();
  auto srcVectorType = srcType.dyn_cast<FVectorType>();
  if (destVectorType && srcVectorType)
    return areTypesWeaklyEquivalent(destVectorType.getElementType(),
                                    srcVectorType.getElementType(), destFlip,
                                    srcFlip);

  // Bundle types are weakly equivalent if all common elements are weakly
  // equivalent.  Non-matching fields are ignored.  Flips are "pushed" into
  // recursive weak type equivalence checks.
  auto destBundleType = destType.dyn_cast<BundleType>();
  auto srcBundleType = srcType.dyn_cast<BundleType>();
  if (destBundleType && srcBundleType)
    return llvm::all_of(destBundleType, [&](auto destElt) -> bool {
      auto destField = destElt.name.getValue();
      auto srcElt = srcBundleType.getElement(destField);
      // If the src doesn't contain the destination's field, that's okay.
      if (!srcElt)
        return true;
      return areTypesWeaklyEquivalent(destElt.type, srcElt->type,
                                      destFlip ^ destElt.isFlip,
                                      srcFlip ^ srcElt->isFlip);
    });

  // Ground types can be connected if their passive, widthless versions
  // are equal and leaf flippedness matches.
  return destType.getWidthlessType() == srcType.getWidthlessType() &&
         destFlip == srcFlip;
}

/// Returns true if the destination is at least as wide as an equivalent source.
bool firrtl::isTypeLarger(FIRRTLBaseType dstType, FIRRTLBaseType srcType) {
  return TypeSwitch<FIRRTLBaseType, bool>(dstType)
      .Case<BundleType>([&](auto dstBundle) {
        auto srcBundle = srcType.cast<BundleType>();
        for (size_t i = 0, n = dstBundle.getNumElements(); i < n; ++i) {
          auto srcElem = srcBundle.getElement(i);
          auto dstElem = dstBundle.getElement(i);
          if (dstElem.isFlip) {
            if (!isTypeLarger(srcElem.type, dstElem.type))
              return false;
          } else {
            if (!isTypeLarger(dstElem.type, srcElem.type))
              return false;
          }
        }
        return true;
      })
      .Case<FVectorType>([&](auto vector) {
        return isTypeLarger(vector.getElementType(),
                            srcType.cast<FVectorType>().getElementType());
      })
      .Default([&](auto dstGround) {
        int32_t destWidth = dstType.getPassiveType().getBitWidthOrSentinel();
        int32_t srcWidth = srcType.getPassiveType().getBitWidthOrSentinel();
        return destWidth <= -1 || srcWidth <= -1 || destWidth >= srcWidth;
      });
}

/// Return the element of an array type or null.  This strips flip types.
Type firrtl::getVectorElementType(Type array) {
  auto vectorType = array.dyn_cast<FVectorType>();
  if (!vectorType)
    return Type();
  return vectorType.getElementType();
}

/// Return the passive version of a firrtl type
/// top level for ODS constraint usage
Type firrtl::getPassiveType(Type anyBaseFIRRTLType) {
  return anyBaseFIRRTLType.cast<FIRRTLBaseType>().getPassiveType();
}

//===----------------------------------------------------------------------===//
// IntType
//===----------------------------------------------------------------------===//

/// Return a SIntType or UInt type with the specified signedness and width.
IntType IntType::get(MLIRContext *context, bool isSigned,
                     int32_t widthOrSentinel) {
  if (isSigned)
    return SIntType::get(context, widthOrSentinel);
  return UIntType::get(context, widthOrSentinel);
}

int32_t IntType::getWidthOrSentinel() {
  if (isa<SIntType>())
    return this->cast<SIntType>().getWidthOrSentinel();
  if (isa<UIntType>())
    return this->cast<UIntType>().getWidthOrSentinel();
  return -1;
}

//===----------------------------------------------------------------------===//
// SIntType
//===----------------------------------------------------------------------===//

SIntType SIntType::get(MLIRContext *context) { return get(context, -1); }

SIntType SIntType::get(MLIRContext *context, std::optional<int32_t> width) {
  if (!width)
    return get(context);
  return get(context, *width);
}

LogicalResult SIntType::verify(function_ref<InFlightDiagnostic()> emitError,
                               int32_t widthOrSentinel) {
  if (widthOrSentinel < -1)
    return emitError() << "invalid width";
  return success();
}

//===----------------------------------------------------------------------===//
// UIntType
//===----------------------------------------------------------------------===//

UIntType UIntType::get(MLIRContext *context) { return get(context, -1); }

UIntType UIntType::get(MLIRContext *context, std::optional<int32_t> width) {
  if (!width)
    return get(context);
  return get(context, *width);
}

LogicalResult UIntType::verify(function_ref<InFlightDiagnostic()> emitError,
                               int32_t widthOrSentinel) {
  if (widthOrSentinel < -1)
    return emitError() << "invalid width";
  return success();
}

//===----------------------------------------------------------------------===//
// Bundle Type
//===----------------------------------------------------------------------===//

struct circt::firrtl::detail::BundleTypeStorage : mlir::TypeStorage {
  using KeyTy = ArrayRef<BundleType::BundleElement>;

  BundleTypeStorage(KeyTy elements)
      : elements(elements.begin(), elements.end()) {
    RecursiveTypeProperties props{true, false, false};
    uint64_t fieldID = 0;
    fieldIDs.reserve(elements.size());
    for (auto &element : elements) {
      auto type = element.type;
      auto eltInfo = type.getRecursiveTypeProperties();
      props.isPassive &= eltInfo.isPassive & !element.isFlip;
      props.containsAnalog |= eltInfo.containsAnalog;
      props.hasUninferredWidth |= eltInfo.hasUninferredWidth;
      fieldID += 1;
      fieldIDs.push_back(fieldID);
      // Increment the field ID for the next field by the number of subfields.
      fieldID += type.getMaxFieldID();
    }
    maxFieldID = fieldID;
    passiveContainsAnalogTypeInfo.setInt(props.toFlags());
  }

  bool operator==(const KeyTy &key) const { return key == KeyTy(elements); }

  static llvm::hash_code hashKey(const KeyTy &key) {
    return llvm::hash_combine_range(key.begin(), key.end());
  }

  static BundleTypeStorage *construct(TypeStorageAllocator &allocator,
                                      KeyTy key) {
    return new (allocator.allocate<BundleTypeStorage>()) BundleTypeStorage(key);
  }

  SmallVector<BundleType::BundleElement, 4> elements;
  SmallVector<uint64_t, 4> fieldIDs;
  uint64_t maxFieldID;

  /// This holds the bits for the type's recursive properties, and can hold a
  /// pointer to a passive version of the type.
  llvm::PointerIntPair<Type, RecursiveTypeProperties::numBits, unsigned>
      passiveContainsAnalogTypeInfo;
};

auto BundleType::getElements() const -> ArrayRef<BundleElement> {
  return getImpl()->elements;
}

/// Return a pair with the 'isPassive' and 'containsAnalog' bits.
RecursiveTypeProperties BundleType::getRecursiveTypeProperties() {
  auto flags = getImpl()->passiveContainsAnalogTypeInfo.getInt();
  return RecursiveTypeProperties::fromFlags(flags);
}

/// Return this type with any flip types recursively removed from itself.
FIRRTLBaseType BundleType::getPassiveType() {
  auto *impl = getImpl();

  // If we've already determined and cached the passive type, use it.
  if (auto passiveType = impl->passiveContainsAnalogTypeInfo.getPointer())
    return passiveType.cast<FIRRTLBaseType>();

  // If this type is already passive, use it and remember for next time.
  if (impl->passiveContainsAnalogTypeInfo.getInt() & IsPassiveBitMask) {
    impl->passiveContainsAnalogTypeInfo.setPointer(*this);
    return *this;
  }

  // Otherwise at least one element is non-passive, rebuild a passive version.
  SmallVector<BundleType::BundleElement, 16> newElements;
  newElements.reserve(impl->elements.size());
  for (auto &elt : impl->elements) {
    newElements.push_back({elt.name, false, elt.type.getPassiveType()});
  }

  auto passiveType = BundleType::get(getContext(), newElements);
  impl->passiveContainsAnalogTypeInfo.setPointer(passiveType);
  return passiveType;
}

std::optional<unsigned> BundleType::getElementIndex(StringAttr name) {
  for (const auto &it : llvm::enumerate(getElements())) {
    auto element = it.value();
    if (element.name == name) {
      return unsigned(it.index());
    }
  }
  return std::nullopt;
}

std::optional<unsigned> BundleType::getElementIndex(StringRef name) {
  for (const auto &it : llvm::enumerate(getElements())) {
    auto element = it.value();
    if (element.name.getValue() == name) {
      return unsigned(it.index());
    }
  }
  return std::nullopt;
}

StringRef BundleType::getElementName(size_t index) {
  assert(index < getNumElements() &&
         "index must be less than number of fields in bundle");
  return getElements()[index].name.getValue();
}

std::optional<BundleType::BundleElement>
BundleType::getElement(StringAttr name) {
  if (auto maybeIndex = getElementIndex(name))
    return getElements()[*maybeIndex];
  return std::nullopt;
}

std::optional<BundleType::BundleElement>
BundleType::getElement(StringRef name) {
  if (auto maybeIndex = getElementIndex(name))
    return getElements()[*maybeIndex];
  return std::nullopt;
}

/// Look up an element by index.
BundleType::BundleElement BundleType::getElement(size_t index) {
  assert(index < getNumElements() &&
         "index must be less than number of fields in bundle");
  return getElements()[index];
}

FIRRTLBaseType BundleType::getElementType(StringAttr name) {
  auto element = getElement(name);
  return element ? element->type : FIRRTLBaseType();
}

FIRRTLBaseType BundleType::getElementType(StringRef name) {
  auto element = getElement(name);
  return element ? element->type : FIRRTLBaseType();
}

FIRRTLBaseType BundleType::getElementType(size_t index) {
  assert(index < getNumElements() &&
         "index must be less than number of fields in bundle");
  return getElements()[index].type;
}

uint64_t BundleType::getFieldID(uint64_t index) {
  return getImpl()->fieldIDs[index];
}

uint64_t BundleType::getIndexForFieldID(uint64_t fieldID) {
  assert(getElements().size() && "Bundle must have >0 fields");
  auto fieldIDs = getImpl()->fieldIDs;
  auto *it = std::prev(llvm::upper_bound(fieldIDs, fieldID));
  return std::distance(fieldIDs.begin(), it);
}

std::pair<FIRRTLBaseType, uint64_t>
BundleType::getSubTypeByFieldID(uint64_t fieldID) {
  if (fieldID == 0)
    return {*this, 0};
  auto fieldIDs = getImpl()->fieldIDs;
  auto subfieldIndex = getIndexForFieldID(fieldID);
  auto subfieldType = getElementType(subfieldIndex);
  auto subfieldID = fieldID - getFieldID(subfieldIndex);
  return {subfieldType, subfieldID};
}

uint64_t BundleType::getMaxFieldID() { return getImpl()->maxFieldID; }

std::pair<uint64_t, bool> BundleType::rootChildFieldID(uint64_t fieldID,
                                                       uint64_t index) {
  auto childRoot = getFieldID(index);
  auto rangeEnd = index + 1 >= getNumElements() ? getMaxFieldID()
                                                : (getFieldID(index + 1) - 1);
  return std::make_pair(fieldID - childRoot,
                        fieldID >= childRoot && fieldID <= rangeEnd);
}

//===----------------------------------------------------------------------===//
// FVectorType
//===----------------------------------------------------------------------===//

struct circt::firrtl::detail::FVectorTypeStorage : mlir::TypeStorage {
  using KeyTy = std::pair<FIRRTLBaseType, size_t>;

  FVectorTypeStorage(KeyTy value) : value(value) {
    auto properties = value.first.getRecursiveTypeProperties();
    passiveContainsAnalogTypeInfo.setInt(properties.toFlags());
  }

  bool operator==(const KeyTy &key) const { return key == value; }

  static FVectorTypeStorage *construct(TypeStorageAllocator &allocator,
                                       KeyTy key) {
    return new (allocator.allocate<FVectorTypeStorage>())
        FVectorTypeStorage(key);
  }

  KeyTy value;

  /// This holds the bits for the type's recursive properties, and can hold a
  /// pointer to a passive version of the type.
  llvm::PointerIntPair<Type, RecursiveTypeProperties::numBits, size_t>
      passiveContainsAnalogTypeInfo;
};

FVectorType FVectorType::get(FIRRTLBaseType elementType, size_t numElements) {
  return Base::get(elementType.getContext(),
                   std::make_pair(elementType, numElements));
}

FIRRTLBaseType FVectorType::getElementType() const {
  return getImpl()->value.first;
}

size_t FVectorType::getNumElements() const { return getImpl()->value.second; }

/// Return the recursive properties of the type.
RecursiveTypeProperties FVectorType::getRecursiveTypeProperties() {
  auto flags = getImpl()->passiveContainsAnalogTypeInfo.getInt();
  return RecursiveTypeProperties::fromFlags(flags);
}

/// Return this type with any flip types recursively removed from itself.
FIRRTLBaseType FVectorType::getPassiveType() {
  auto *impl = getImpl();

  // If we've already determined and cached the passive type, use it.
  if (auto passiveType = impl->passiveContainsAnalogTypeInfo.getPointer())
    return passiveType.cast<FIRRTLBaseType>();

  // If this type is already passive, return it and remember for next time.
  if (impl->passiveContainsAnalogTypeInfo.getInt() & IsPassiveBitMask) {
    impl->passiveContainsAnalogTypeInfo.setPointer(*this);
    return *this;
  }

  // Otherwise, rebuild a passive version.
  auto passiveType =
      FVectorType::get(getElementType().getPassiveType(), getNumElements());
  impl->passiveContainsAnalogTypeInfo.setPointer(passiveType);
  return passiveType;
}

uint64_t FVectorType::getFieldID(uint64_t index) {
  return 1 + index * (getElementType().getMaxFieldID() + 1);
}

uint64_t FVectorType::getIndexForFieldID(uint64_t fieldID) {
  assert(fieldID && "fieldID must be at least 1");
  // Divide the field ID by the number of fieldID's per element.
  return (fieldID - 1) / (getElementType().getMaxFieldID() + 1);
}

std::pair<FIRRTLBaseType, uint64_t>
FVectorType::getSubTypeByFieldID(uint64_t fieldID) {
  if (fieldID == 0)
    return {*this, 0};
  auto subfieldIndex = getIndexForFieldID(fieldID);
  return {getElementType(), fieldID - getFieldID(subfieldIndex)};
}

uint64_t FVectorType::getMaxFieldID() {
  return getNumElements() * (getElementType().getMaxFieldID() + 1);
}

std::pair<uint64_t, bool> FVectorType::rootChildFieldID(uint64_t fieldID,
                                                        uint64_t index) {
  auto childRoot = getFieldID(index);
  auto rangeEnd =
      index >= getNumElements() ? getMaxFieldID() : (getFieldID(index + 1) - 1);
  return std::make_pair(fieldID - childRoot,
                        fieldID >= childRoot && fieldID <= rangeEnd);
}

//===----------------------------------------------------------------------===//
// RefType
//===----------------------------------------------------------------------===//

namespace circt {
namespace firrtl {
namespace detail {
struct RefTypeStorage : mlir::TypeStorage {
  using KeyTy = FIRRTLBaseType;

  RefTypeStorage(KeyTy value) : value(value) {}

  bool operator==(const KeyTy &key) const { return key == value; }

  static RefTypeStorage *construct(TypeStorageAllocator &allocator, KeyTy key) {
    return new (allocator.allocate<RefTypeStorage>()) RefTypeStorage(key);
  }

  KeyTy value;
};

} // namespace detail
} // namespace firrtl
} // namespace circt

auto RefType::get(FIRRTLBaseType type) -> RefType {
  return Base::get(type.getContext(), type);
}

auto RefType::getType() -> FIRRTLBaseType { return getImpl()->value; }

auto RefType::verify(function_ref<InFlightDiagnostic()> emitErrorFn,
                     FIRRTLBaseType base) -> LogicalResult {
  if (!base.isPassive())
    return emitErrorFn() << "reference base type must be passive";
  return success();
}

//===----------------------------------------------------------------------===//
// AnalogType
//===----------------------------------------------------------------------===//

AnalogType AnalogType::get(mlir::MLIRContext *context) {
  return AnalogType::get(context, -1);
}

AnalogType AnalogType::get(mlir::MLIRContext *context,
                           std::optional<int32_t> width) {
  if (!width)
    return AnalogType::get(context);
  return AnalogType::get(context, *width);
}

LogicalResult AnalogType::verify(function_ref<InFlightDiagnostic()> emitError,
                                 int32_t widthOrSentinel) {
  if (widthOrSentinel < -1)
    return emitError() << "invalid width";
  return success();
}

//===----------------------------------------------------------------------===//
// FIRRTLDialect
//===----------------------------------------------------------------------===//

void FIRRTLDialect::registerTypes() {
  addTypes<SIntType, UIntType, ClockType, ResetType, AsyncResetType, AnalogType,
           // Derived Types
           BundleType, FVectorType, RefType>();
}

// Get the bit width for this type, return None  if unknown. Unlike
// getBitWidthOrSentinel(), this can recursively compute the bitwidth of
// aggregate types. For bundle and vectors, recursively get the width of each
// field element and return the total bit width of the aggregate type. This
// returns None, if any of the bundle fields is a flip type, or ground type with
// unknown bit width.
std::optional<int64_t> firrtl::getBitWidth(FIRRTLBaseType type,
                                           bool ignoreFlip) {
  std::function<std::optional<int64_t>(FIRRTLBaseType)> getWidth =
      [&](FIRRTLBaseType type) -> std::optional<int64_t> {
    return TypeSwitch<FIRRTLBaseType, std::optional<int64_t>>(type)
        .Case<BundleType>([&](BundleType bundle) -> std::optional<int64_t> {
          int64_t width = 0;
          for (auto &elt : bundle) {
            if (elt.isFlip && !ignoreFlip)
              return std::nullopt;
            auto w = getBitWidth(elt.type);
            if (!w.has_value())
              return std::nullopt;
            width += *w;
          }
          return width;
        })
        .Case<FVectorType>([&](auto vector) -> std::optional<int64_t> {
          auto w = getBitWidth(vector.getElementType());
          if (!w.has_value())
            return std::nullopt;
          return *w * vector.getNumElements();
        })
        .Case<IntType>([&](IntType iType) { return iType.getWidth(); })
        .Case<ClockType, ResetType, AsyncResetType>([](Type) { return 1; })
        .Default([&](auto t) { return std::nullopt; });
  };
  return getWidth(type);
}
